[python] Add direct LeRobot capture writer - #9665
Conversation
14eea28 to
c7b0a9d
Compare
Add a LeRobot-compatible writer that batches completed episodes into multimodal Paimon commits, resumes compatible tables, and stores image frames as PNG BLOBs. Co-Authored-By: Codex <noreply@anthropic.com> AI-Model: gpt-5 AI-Contributed/Feature: 424/424 AI-Contributed/UT: 366/366
Allow table commits to attach snapshot properties and use them to resume LeRobot frame, episode, and task indices without scanning table data. Co-Authored-By: Codex <noreply@anthropic.com> AI-Model: gpt-5 Co-Authored-By: Codex <noreply@ai-tool.com> Co-Authored-By: Codex <noreply@openai.com> AI-Contributed/Feature: 71/71 AI-Contributed/UT: 40/40
b4aedae to
0c3ebdd
Compare
JingsongLi
left a comment
There was a problem hiding this comment.
Found two reproducible input-compatibility issues in the direct capture writer; details are inline.
| return _encode_media_frame( | ||
| value, | ||
| channel_first=actual_shape == expected_shape, |
There was a problem hiding this comment.
[P1] Preserve the layout of native LeRobot HWC camera features
Matching the declared shape does not imply CHW. LeRobot 0.4.4's hw_to_dataset_features preserves the camera's HWC shape and sets names=["height", "width", "channels"]. With a standard image feature declaring shape=(480, 640, 3) and a matching np.uint8 frame, this sets channel_first=True and transposes the image to (640, 3, 480). I reproduced the first add_frame() failing with Unsupported LeRobot media frame shape or dtype: (640, 3, 480), uint8, so the advertised native recording-loop integration cannot record standard camera frames. Please determine the layout from the feature's dimension names and actual channel axis, and add a regression test using native HWC feature metadata.
| raise ValueError( | ||
| "LeRobot feature %s expected shape %s, got %s." | ||
| % (name, expected_shape, value.shape)) | ||
| return _normalize_value(value, feature, name) |
There was a problem hiding this comment.
[P2] Normalize validated NumPy arrays before passing them to the Arrow helpers
For non-scalar features, _normalize_value() returns the original ndarray, whereas the existing import path supplies Python values from to_pylist(). This makes valid native LeRobot inputs fail in the subsequent _safe_array([value], ...): a float32 array with shape (2, 2) raises Can only convert 1-dimensional array values, and a bool array with shape (2,) is rejected because its np.bool_ elements are not Python bool instances. Both satisfy LeRobot's native dtype/shape validation, and I reproduced both failures through PaimonLeRobotWriter.add_frame(). Please convert validated arrays to Python scalars/nested lists before buffering and Arrow validation, and cover both cases. Applying .tolist() after validation allowed both examples to complete a write/read round trip.
Summary
Add a LeRobot recording-loop compatible writer that sends captured frames directly to a Paimon multimodal table without materializing an intermediate LeRobot dataset.
This PR builds on the generic PyPaimon Snapshot-properties commit support merged in #9692 and uses it to persist LeRobot resume state.
Changes
PaimonLeRobotWriterwith the recording lifecycle used by LeRobot: frame buffering, episode save/discard, pending-frame checks, and finalize.task, generated indices excluded, exact numeric dtype/shape, and CHW/HWC image support.finalize(); optional positiveepisodes_per_commitandflush()provide explicit intermediate commit boundaries.clear_episode_buffer()limited to the current unsaved episode; already accepted episodes remain in the pending batch.Testing
python -m pytest -q pypaimon/tests/multimodal_lerobot_test.py pypaimon/tests/multimodal_lerobot_writer_test.py pypaimon/tests/table_commit_test.py pypaimon/tests/table/simple_table_test.py(92 passed)git diff --checkNotes
This draft intentionally starts with PNG image frames, synchronous encoding, and a single active writer per table. MP4 encoding, concurrent index allocation, and streaming-checkpoint semantics are deferred.